package com.tianji.aigc.agent;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constants;
import com.tianji.aigc.domain.vo.ChatEventVO;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.common.utils.UserContext;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.tianji.aigc.constants.Constants.*;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

/**
 * @author Administrator
 * @version 1.0
 * @description: TODO
 * @date 2025/7/24 12:26
 */
public abstract class AbstractAgent implements Agent {
    @Autowired
    private ChatClient chatClient;
    @Autowired
    private ChatMemory chatMemory;
    @Autowired
    private IChatSessionService sessionService;
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();
    public static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    /**
     * 1- 普通聊天 （Day3）
     *
     * @param question
     * @param sessionId
     * @return
     */
    @Override
    public String chat(String question, String sessionId) {
        // 1- 生成requestId，用于唯一标识本次请求，供后续课程Tool中的容器使用（课程卡片效果）
        String requestId = IdUtil.fastSimpleUUID();

        // 2-获取用户ID，放入工具上下文中，供后续预下单工具使用（预下单）
        Long userId = UserContext.getUser();

        // 3- 异步更新会话标题
        sessionService.update(sessionId,userId,question);

        return this.getChatClientRequest(question, sessionId, requestId)  // 1- 用户提示词
                .call()
                .content();
    }

    /**
     * 2、流式聊天 （Day3）
     *
     * @param question
     * @param sessionId
     * @return
     */
    @Override
    public Flux<ChatEventVO> streamChat(String question, String sessionId) {
        // 1- 大模型输出内容的缓存器，用于在输出中断后的数据存储
        StringBuilder sessionRecord = new StringBuilder();

        // 2- 生成requestId，用于唯一标识本次请求，供后续课程Tool中的容器使用（课程卡片效果）
        String requestId = IdUtil.fastSimpleUUID();

        // 3-获取用户ID，放入工具上下文中，供后续预下单工具使用（预下单）
        Long userId = UserContext.getUser();

        // 4- 异步更新会话标题
        sessionService.update(sessionId,userId,question);

        return this.getChatClientRequest(question, sessionId, requestId)
                .stream() // 1- 流式聊天
                .chatResponse() // 1- 获取流式聊天结果

                .doFirst(() -> GENERATE_STATUS.put(sessionId, true))  // 3- 在处理大模型生成元素前，将输出标记置为true
                .doOnComplete(() -> GENERATE_STATUS.remove(sessionId))//  3- 在处理大模型生成完成后，删除标记
                .doOnError(throwable -> GENERATE_STATUS.remove(sessionId))  // 3- 在大模型内容出异常时，删除标记
                .doOnCancel(() -> addStopSessionRecord(sessionId, sessionRecord.toString())) // 5- 当输出被取消时(比如takeWhile判断为false时)，保存输出的内容到历史记录中，【注意】：一定要放在takeWhile前面，否则不生效
                .takeWhile(p -> GENERATE_STATUS.getOrDefault(sessionId, false))//3- 在处理元素前，判断生成标记是否为true，为true则继续处理，false结束流处理

                .map(chatResponse -> {  // 1- 处理流式聊天结果 -> 封装成ChatEventVO
                    // 1、在大模型最后一次输出时，建立 requestId 与 messageId的映射关系，供后续会话记忆存储器使使用（用于解决卡片记忆丢失Bug）
                    String finishReason = chatResponse.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(STOP, finishReason)) {
                        String messageId = chatResponse.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constants.REQUEST_ID, requestId);
                    }

                    // 2、获取大模型输出的内容
                    String text = chatResponse.getResult().getOutput().getText(); // 1

                    // 3、将大模型回复添加至当前会话记录中
                    sessionRecord.append(text);  // 5- 追加到当前会话记录中

                    // 4、封装成ChatEventVO
                    return ChatEventVO.builder()  // 1
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                // 在流响应最后，添加卡片数据 和 流结束标识数据
                .concatWith(Flux.defer(() -> {
                    // 通过请求id获取工具中放入的卡片数据，如果不为空，将其追加到流最后
                    Map<String, Object> map = ToolResultHolder.get(requestId);
                    if (ObjectUtil.isNotEmpty(map)) {
                        ToolResultHolder.remove(requestId);

                        // 封装卡片数据（ChatEventVO）
                        ChatEventVO chatEventVo = ChatEventVO.builder()  // 1
                                .eventData(map)
                                .eventType(ChatEventTypeEnum.PARAM.getValue())
                                .build();
                        return Flux.just(chatEventVo, STOP_EVENT);
                    }
                    return Flux.just(STOP_EVENT);
                }));
    }

    private ChatClient.ChatClientRequestSpec getChatClientRequest(String question, String sessionId, String requestId) {
        return chatClient.prompt()
                .system(promptSystem -> promptSystem.text(this.systemMessage()).params(this.systemMessageParams()))
                .advisors(advisor -> advisor.params(this.advisorParams(sessionId)).advisors(this.advisors()))
                .tools(this.tools())
                .toolContext(this.toolContext(requestId))
                .user(question);
    }

    /**
     * 2、停止对话 （Day1 - 随堂4）
     *
     * @param sessionId@return
     */
    @Override
    public void stop(String sessionId) {
        GENERATE_STATUS.remove(sessionId); // 3- 停止生成，删除标记
    }

    /**
     * 创建并返回一个Advisor的Map
     * @return
     */
    @Override
    public Map<String, Object> advisorParams(String sessionId){
        return Map.of(CHAT_MEMORY_CONVERSATION_ID_KEY, sessionId);
    };

    /**
     * 3、添加打断时的会话记录
     *
     * @param sessionId
     * @param hisRecord
     */
    private void addStopSessionRecord(String sessionId, String hisRecord) {
        chatMemory.add(sessionId, new AssistantMessage(hisRecord));  // 5- 添加打断时的会话记录
    }
}
